import mindspore as ms
import mindspore.nn as nn
from util.datasets import load_mnist
from util.transform import data_iter
from util.callback import AccuracyMonitor
# from mindvision.classification.dataset import Mnist


def main():
    batch_size, lr, num_epochs = 256, 0.1, 10

    # dataset_train = Mnist(path='/shareData/mindspore-dataset/Mnist', split="train",
    #                       batch_size=batch_size, repeat_num=1, shuffle=True,
    #                       download=False).run()
    # dataset_test = Mnist(path='/shareData/mindspore-dataset/Mnist', split="test",
    #                      batch_size=batch_size, repeat_num=1, shuffle=True,
    #                      download=False).run()

    # net = nn.SequentialCell(
    #     nn.Flatten(),
    #     nn.Dense(32*32, 256),
    #     nn.ReLU(),
    #     nn.Dense(256, 10),
    # )

    features_t, labels_t = load_mnist('/shareData/mindspore-dataset/Mnist/train')
    features_v, labels_v = load_mnist('/shareData/mindspore-dataset/Mnist/test', split='t10k')

    dataset_train = data_iter(features_t, labels_t, batch_size)
    dataset_valid = data_iter(features_v, labels_v, batch_size)

    net = nn.SequentialCell(
        nn.Dense(28 * 28, 256),
        nn.ReLU(),
        nn.Dense(256, 10),
    )

    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    opti = nn.SGD(net.trainable_params(), learning_rate=lr)

    model = ms.Model(net, loss_fn=loss, optimizer=opti, metrics={'acc'})

    model.train(num_epochs, dataset_train, callbacks=[AccuracyMonitor(dataset_valid)])


if __name__ == '__main__':
    main()


"""
output
epoch:[1/10] Loss:0.44541082 Train Accuracy:0.8763166666666666 Valid Accuracy:0.8811
epoch:[2/10] Loss:0.34548318 Train Accuracy:0.9013166666666667 Valid Accuracy:0.9055
epoch:[3/10] Loss:0.2998726 Train Accuracy:0.9152166666666667 Valid Accuracy:0.9186
epoch:[4/10] Loss:0.27463746 Train Accuracy:0.9216166666666666 Valid Accuracy:0.9249
epoch:[5/10] Loss:0.2511615 Train Accuracy:0.9294 Valid Accuracy:0.9322
epoch:[6/10] Loss:0.23027128 Train Accuracy:0.9350166666666667 Valid Accuracy:0.9362
epoch:[7/10] Loss:0.21230069 Train Accuracy:0.9410833333333334 Valid Accuracy:0.9415
epoch:[8/10] Loss:0.1968191 Train Accuracy:0.94505 Valid Accuracy:0.9437
epoch:[9/10] Loss:0.1849168 Train Accuracy:0.9482 Valid Accuracy:0.9457
epoch:[10/10] Loss:0.1722716 Train Accuracy:0.95185 Valid Accuracy:0.9498
"""